from geometry_process import *
from multiprocess import Pool
import random
from deap import base, creator, tools, algorithms
import os
import copy
import dill as pickle

# 生成实例的代码，返回一个列表module_list
def my_indices(res_dict):
    print(f"Creating individual in process {os.getpid()}")
    module_list = []
    # 生成与初始blocks数量相当的modules
    module_list = [Module() for _ in range(len(res_dict))]
    # 打乱原始的block list，使得初始种群尽可能多样化
    shuffled_block_list = np.random.permutation(len(res_dict))
    for each_block in shuffled_block_list:
        for each_module in module_list:
            # 将每一个block塞入第一个可行的module中
            each_module.add_block([each_block])
            if each_module.is_feasible():
                break
            else:
                each_module.delete_block([each_block])
    return module_list

# 估计实例fitness的函数
def eval_func(individual):
    #print(f"Evaluating in process {os.getpid()}")
    module_num = 0
    diff_module_num = 0
    small_module_num = 0
    dimensions = []

    for module in individual:
        #if not module.is_feasible():
            #return (1e9, 1e9,)
        if len(module.block) != 0:
            module_num += 1

            # 检查同种模块数量
            # 获取模块的长和宽，并确保长和宽是排序过的
            dimension = get_width_length(module.get_geoms())
            if dimension not in dimensions:
                diff_module_num += 1
                dimensions.append(dimension)
            
            if max(dimension) <= 1.5:
                small_module_num += 1
    return (module_num, diff_module_num,)

# 两点交叉函数，用于交配生成子代
def my_cxTwoPoint(ind1, ind2):
    size = min(len(ind1), len(ind2))
    cxpoint1 = random.randint(1, size)
    cxpoint2 = random.randint(1, size - 1)
    if cxpoint2 >= cxpoint1:
        cxpoint2 += 1
    else:  # Swap the two cx points
        cxpoint1, cxpoint2 = cxpoint2, cxpoint1
    
    ind1[cxpoint1:cxpoint2], ind2[cxpoint1:cxpoint2] = ind2[cxpoint1:cxpoint2], ind1[cxpoint1:cxpoint2]
    
    #ind1 = refine_gene(ind1, ind1[cxpoint1:cxpoint2])
    #ind2 = refine_gene(ind2, ind2[cxpoint1:cxpoint2])
    return ind1, ind2,

# 交叉函数，随即单点交叉
def my_cxDiscontinuous(ind1, ind2):
    size = min(len(ind1), len(ind2))
    num_genes = random.randint(1, size)  # 随机确定交换的基因数量
    cxpoints = random.sample(range(size), num_genes)  # 随机选择一组交叉点

    # 在交叉点上分别交换基因
    for cxpoint in cxpoints:
        ind1[cxpoint], ind2[cxpoint] = ind2[cxpoint], ind1[cxpoint]

    return ind1, ind2

def refine_gene(ind):
    blocks = set()
    new_ind = []

    # 删除重复的blocks
    for each in ind:
        ind_blocks = set(each.block)
        if not ind_blocks & blocks:
            new_ind.append(each)
            blocks.update(ind_blocks)

    while len(new_ind) < len(ind):
        new_ind.append(Module())

    # 填充缺失的blocks
    miss_blocks = set(res_dict.keys()) - blocks
    for each_block in miss_blocks:
        # 找到第一个空模块
        first_empty_module = next((m for m in new_ind if not m.block), None)
        added = False
        for each_module in new_ind:
            if each_module.block:  # 非空模块
                each_module.add_block([each_block])
                if each_module.is_feasible():
                    added = True
                    break
                else:
                    each_module.delete_block([each_block])
        if not added and first_empty_module:  # 加入到第一个空模块
            first_empty_module.add_block([each_block])
            

    return creator.Individual(new_ind)

# 变异函数
def my_mut(ind, indpb):

    new_ind = []
    temp_ind = []
    # 随机删除一些modules
    for i in range(len(ind)):
        if random.random() >= indpb:
            new_ind.append(ind[i])
            temp_ind.append(copy.deepcopy(ind[i]))
    # 将空modules重新加入list以保证长度相等
    while len(new_ind) < len(ind):
        new_ind.append(Module())
        temp_ind.append(Module())
    return creator.Individual(new_ind),

class Module():

    def __init__(self):
        self.block = []
    
    # 返回block list，list的元素为res_dict中的block id
    def get_block(self):
        result = [res_dict[key] for key in self.block if key in res_dict]
        return result
    
    # 返回是否该Module为空
    def is_empty(self):
        return len(self.block) == 0

    # 返回是否该Module为矩形。如果Module为空，则也返回True
    def is_rectangle(self):
        blocks = self.get_block()
        if len(blocks) == 0:
            return True

        geom = self.get_geoms()
        if isinstance(geom, Polygon):
            l = geom.length
            a = geom.area
            minx, miny, maxx, maxy = geom.bounds
            c_l = (maxx - minx + maxy - miny) * 2
            c_a = (maxx - minx) * (maxy - miny)
            return abs(l - c_l) < 0.1 and abs(a - c_a) < 0.1
        else:
            return False
    
    # 返回是否该Module未超限。如果未超限且为矩形，返回True
    def is_feasible(self):
        if self.is_rectangle():
            if len(self.get_block()) == 0:
                return True
            geom = self.get_geoms()
            width, length = get_width_length(geom)
            return (width <= max_width and length <= max_length) or (width <= max_length and length <= max_width)
        else:
            return False

    # 向block list中添加block ids，参数为一个列表
    def add_block(self, ind_list):
        self.block += ind_list

    # 删除block list中的block ids
    def delete_block(self, ind_list):
        for each in ind_list:
            self.block.remove(each)
    
    # 将block list中的blocks合并后返回合并图形。
    def get_geoms(self):
        blocks = self.get_block()
        return unary_union(blocks)
    
    # 返回Module的面积
    def get_area(self):
        geom = self.get_geoms()
        return geom.area

def generate_individual(res_dict):
    """这个函数用于生成单个个体，可以将res_dict作为参数传递给my_indices函数"""
    return creator.Individual(my_indices(res_dict))

max_width = 2.5
max_length = 8

case = 'case2'
with open('checkpoints/'+case+'/seg.pkl', 'rb') as seg_file:
    seg = pickle.load(seg_file)

res_dict = {}
for i,poly in enumerate(list(seg.geoms)):
    res_dict[i] = poly

# 创建类型
creator.create("FitnessMin", base.Fitness, weights=(-0.3, -1.0))
creator.create("Individual", list, fitness=creator.FitnessMin)

# 注册函数
toolbox = base.Toolbox()

# 注册生成单个实例的函数，实例类型为creator.Individual，生成函数为lambda: my_indices(res_dict)
toolbox.register("individual", tools.initIterate, creator.Individual, lambda: my_indices(res_dict))
# 注册生成种群的函数
toolbox.register("population", tools.initRepeat, list, toolbox.individual)
# 注册交配方法
toolbox.register("mate", my_cxTwoPoint)
# 注册变异方法
toolbox.register("mutate", my_mut, indpb=0.2)
# 注册种群选优方法
toolbox.register("select", tools.selTournament, tournsize=3)
# 注册评估方法
toolbox.register("evaluate", eval_func)

# 遗传算法参数
pop_size = 32
cx_prob = 0.5
mut_prob = 0.5
n_generations = 40
if __name__ == "__main__":
    # 多线程
    pool = Pool(processes=16)
    toolbox.register("map", pool.map)
    # 准备生成种群所需的参数，这里假设res_dict已经在外部定义
    args = [(res_dict,) for _ in range(pop_size)]

    # 使用多进程生成种群
    pop = pool.starmap(generate_individual, args)
    pop_path = 'checkpoints/'+case+'/pop.pkl'
    with open(pop_path, 'wb') as result_file:
        pickle.dump(pop, result_file)

    # 注册统计方法
    stats = tools.Statistics(lambda ind: sum(ind.fitness.values))
    stats.register("avg", np.mean)
    stats.register("min", np.min)
    logbook = []
    for gen in range(n_generations):
            # 选择、交叉、变异...
            offspring = algorithms.varAnd(pop, toolbox, cxpb=cx_prob, mutpb=mut_prob)
            offspring = pool.map(refine_gene, offspring)
            # 评估新一代个体
            fitnesses = toolbox.map(toolbox.evaluate, offspring)
            for ind, fit in zip(offspring, fitnesses):
                ind.fitness.values = fit

            # 更新种群
            pop[:] = toolbox.select(offspring, len(pop))
            # 记录统计数据
            fits = [(ind.fitness.values[0]*0.3+ind.fitness.values[1]*1.0) for ind in pop]
            length = len(pop)
            mean = sum(fits) / length
            minimum = min(fits)
            logbook.append([mean, minimum])
            print(f"Gen: {gen}, Pop: {length}, Avg Fit: {mean:.2f}, Min Fit: {minimum:.2f}")
    #best_individual = tools.selBest(result, k=50)
    pool.close()
    pool.join()
    #print("Fitness: ", best_individual[0].fitness.values)
    #draw_type_multipolygon(best_individual[0])
    result_path = 'checkpoints/'+case+'/result.pkl'
    logbook_path = 'checkpoints/'+case+'/logbook.pkl'

    # 序列化并保存 result
    with open(result_path, 'wb') as result_file:
        pickle.dump(pop, result_file)

    # 序列化并保存 logbook
    with open(logbook_path, 'wb') as logbook_file:
        pickle.dump(logbook, logbook_file)